import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings('ignore')

class Backtest:
    def __init__(self, data, strategy, initial_capital=1e6, top_n=10):
        """
        初始化回测
        :param data: 包含信号和收益率的数据
        :param strategy: 策略对象
        :param initial_capital: 初始资金
        :param top_n: 买入前top_n只股票
        """
        self.data = data.copy()
        self.strategy = strategy
        self.initial_capital = initial_capital
        self.top_n = top_n
        self.results = None

    def run(self):
        """运行回测"""
        dates = sorted(self.data['date'].unique())
        portfolio_values = []
        current_cash = self.initial_capital
        current_stock_value = 0
        current_portfolio_value = current_cash+ current_stock_value
        positions = {}  # {stock: {'shares': quantity, 'entry_price': price}}

        historical_data = pd.DataFrame()  # Accumulate all historical data

        for i, date in enumerate(dates):
            # Get today's market data
            daily_data = self.data[self.data['date'] == date].copy()
            historical_data = pd.concat([historical_data, daily_data])

            # Prepare all information for strategy
            strategy_input = {
                'current_date': date,
                'current_cash': current_cash,
                'current_positions': positions.copy(),
                'historical_data': historical_data.copy(),
                'current_portfolio_value': current_portfolio_value
            }

            # Get trading decisions from strategy
            trades = self.strategy.generate_trades(strategy_input)
            # Execute trades
            for trade in trades:
                stock = trade['stock']
                action = trade['action']  # 'buy' or 'sell'
                shares = trade.get('shares', 0)  # Optional: exact share quantity
                amount = trade.get('amount', 0)  # Optional: cash amount to use

                if stock not in daily_data['stock'].values:
                    continue  # Skip if stock not available today

                stock_price = daily_data[daily_data['stock'] == stock]['close'].values[0]

                if action == 'buy':
                    # Calculate shares to buy
                    if shares <= 0 and amount <= 0:
                        continue  # Invalid trade

                    if amount > 0:
                        shares_to_buy = int(amount / stock_price)
                    else:
                        shares_to_buy = shares

                    # Execute buy
                    cost = shares_to_buy * stock_price
                    if cost > current_cash:
                        continue  # Not enough cash

                    current_cash -= cost
                    if stock in positions:
                        # Average the entry price if adding to existing position
                        total_shares = positions[stock]['shares'] + shares_to_buy
                        positions[stock]['entry_price'] = (positions[stock]['entry_price'] * positions[stock][
                                                            'shares'] + stock_price * shares_to_buy
                                                          ) / total_shares
                        positions[stock]['shares'] = total_shares
                    else:
                        positions[stock] = {
                            'shares': shares_to_buy,
                            'entry_price': stock_price
                        }

                elif action == 'sell':
                    if stock not in positions or positions[stock]['shares'] <= 0:
                        continue  # No position to sell

                    # Calculate shares to sell
                    if shares <= 0:
                        shares_to_sell = positions[stock]['shares']  # Sell all if not specified
                    else:
                        shares_to_sell = min(shares, positions[stock]['shares'])

                    # Execute sell
                    proceeds = shares_to_sell * stock_price
                    current_cash += proceeds
                    positions[stock]['shares'] -= shares_to_sell

                    # Remove position if zero shares
                    if positions[stock]['shares'] <= 0:
                        del positions[stock]

            # Calculate portfolio value after trades
            current_stock_value = 0
            for stock, pos_info in positions.items():
                if stock in daily_data['stock'].values:
                    stock_close = daily_data[daily_data['stock'] == stock]['close'].values[0]
                    current_stock_value += pos_info['shares'] * stock_close
                    print(stock,stock_close,pos_info['shares'])
            print(current_cash,current_stock_value)
            portfolio_value = current_stock_value+current_cash
            portfolio_values.append(portfolio_value)

        # Save backtest results
        self.results = pd.DataFrame({
            'date': dates,
            'portfolio_value': portfolio_values
        })

        return self.results

    def evaluate_performance(self):
        """评估回测表现"""
        if self.results is None:
            raise ValueError("请先运行回测")

        results = self.results.copy()
        results['daily_return'] = results['portfolio_value'].pct_change()

        # 计算年化收益率
        total_return = (results['portfolio_value'].iloc[-1] / self.initial_capital) - 1
        days = len(results)
        annualized_return = (1 + total_return) ** (252 / days) - 1

        # 计算年化波动率
        annualized_vol = results['daily_return'].std() * np.sqrt(252)

        # 计算夏普比率
        sharpe_ratio = annualized_return / annualized_vol if annualized_vol != 0 else np.nan

        # 计算最大回撤
        results['peak'] = results['portfolio_value'].cummax()
        results['drawdown'] = (results['portfolio_value'] - results['peak']) / results['peak']
        max_drawdown = results['drawdown'].min()

        # 计算胜率
        win_rate = (results['daily_return'] > 0).mean()

        return {
            'Total Return': total_return,
            'Annualized Return': annualized_return,
            'Annualized Volatility': annualized_vol,
            'Sharpe Ratio': sharpe_ratio,
            'Max Drawdown': max_drawdown,
            'Win Rate': win_rate,
            'Top N': self.top_n
        }

    def plot_results(self):
        """绘制回测结果"""
        if self.results is None:
            raise ValueError("请先运行回测")

        plt.figure(figsize=(12, 6))
        plt.plot(self.results['date'], self.results['portfolio_value'])
        plt.title('Portfolio Value Over Time')
        plt.xlabel('Date')
        plt.ylabel('Portfolio Value')
        plt.grid(True)
        plt.show()